from typing import Dict, Tuple, Optional, Any
from threading import Lock

from axiom_boot.conf.manager import Settings
from axiom_boot.di.decorators import service
from axiom_boot.di.dependency import autowired
from axiom_boot.vector.factory import PGVectorStoreManager

from .node_postprocessor_interface import NodePostprocessorInterface
from .sentence_transformer_rerank import SentenceTransformerRerankProcessor
from .contextual_expansion_config import ContextualExpansionConfig
from .contextual_expansion_postprocessor import ContextualExpansionPostprocessor


@service()
class NodePostprocessorManager:
    """
    一个管理器服务，负责根据配置按需创建和缓存后处理器实例。
    """

    def __init__(
        self,
        settings: Settings = autowired(),
        vector_store_manager: PGVectorStoreManager = autowired(),
    ):
        self._settings = settings
        self._vector_store_manager = vector_store_manager
        self._postprocessors: Dict[Tuple[str, Optional[int]], NodePostprocessorInterface] = {}
        self._lock = Lock()

    def get_postprocessor(
        self,
        profile_name: str,
        # 新增参数以支持需要向量存储的后处理器
        collection_name: Optional[str] = None,
        embed_dim: Optional[int] = None,
        **kwargs,
    ) -> NodePostprocessorInterface:
        """
        根据配置档案和动态参数获取后处理器实例。
        """
        # [BugFix] 缓存键应该包含所有动态参数，这里用 kwargs 的元组作为键的一部分
        cache_key = (profile_name, collection_name, embed_dim, tuple(sorted(kwargs.items())))
        if cache_key not in self._postprocessors:
            with self._lock:
                if cache_key not in self._postprocessors:
                    self._postprocessors[cache_key] = self._create_postprocessor(
                        profile_name, collection_name, embed_dim, **kwargs
                    )
        
        return self._postprocessors[cache_key]

    def _create_postprocessor(
        self,
        profile_name: str,
        collection_name: Optional[str] = None,
        embed_dim: Optional[int] = None,
        **kwargs,
    ) -> NodePostprocessorInterface:
        """
        根据配置档案和可选的动态 top_n 创建一个新的后处理器服务实例。
        """
        # [BugFix] 实现大小写不敏感的档案查找
        profile_config = None
        profiles_dict = self._settings.vector.postprocessor.profiles
        
        # 遍历配置字典，用不区分大小写的方式查找匹配的key
        for key, config in profiles_dict.items():
            if key.lower() == profile_name.lower():
                profile_config = config
                break

        if not profile_config:
            raise ValueError(f"未在配置中找到名为 '{profile_name}' 的后处理器 (postprocessor) 档案。")

        provider = profile_config.provider
        # 合并配置中的参数和动态传入的参数，动态参数优先
        # [BugFix] 确保 profile_config.config 不为 None
        profile_kwargs = profile_config.config or {}
        final_kwargs = {**profile_kwargs, **kwargs}

        if provider == "sentence-transformer-rerank":
            # [BugFix] 将 model_name 和 top_n 也纳入 final_kwargs
            if profile_config.model_name:
                final_kwargs.setdefault('model_name', profile_config.model_name)
            if profile_config.top_n is not None:
                final_kwargs.setdefault('top_n', profile_config.top_n)

            return SentenceTransformerRerankProcessor(**final_kwargs)
        
        if provider == "contextual-expansion":
            if not all([collection_name, embed_dim]):
                raise ValueError("ContextualExpansionPostprocessor 需要 collection_name 和 embed_dim。")
            
            config = ContextualExpansionConfig(**final_kwargs)
            vector_store = self._vector_store_manager.get_store(collection_name, embed_dim)
            return ContextualExpansionPostprocessor(
                vector_store=vector_store,
                window_size=config.window_size
            )
        
        raise TypeError(f"不支持的后处理器 (postprocessor) provider: {provider}") 